#!/usr/bin/env python

import cv2
import os
import rospy
import threading
from cv_bridge import CvBridge
from sensor_msgs.msg import Image
from interbotix_perception_modules.srv import SnapPicture, SnapPictureResponse

# To start this application, open a terminal on the robot and type...
#   'roslaunch interbotix_perception_modules picture_snapper.launch'
#
# To save a picture, open a new terminal in this directory and type...
#   'rosservice call /apriltag/snap_picture "filename: '[filename].jpg'"'

### @brief Class to get the latest picture from the camera and save it to a file with a specified name
### @details - meant to run on the same computer with the AprilTag Single Image Server node;
###            this way, all images will be saved to the same computer so that the AprilTag
###            Single Image Server node can access them
class PictureSnapper(object):
    def __init__(self):
        self.image = None
        self.img_mutex = threading.Lock()

        # topic containing raw rgb data
        self.camera_color_topic = rospy.get_param("camera_color_topic").strip("/")
        self.sub_camera_color = rospy.Subscriber("/" + self.camera_color_topic, Image, self.camera_color_cb)

        apriltag_ns = rospy.get_param("~apriltag_ns", default="apriltag")

        # directory in which to save pictures (leave this as is)
        dir_param = rospy.get_param(
            "/" + apriltag_ns + "/picture_snapper/image_save_dir", 
            default="interbotix/picture_snapper/")
        self.image_save_dir = os.path.dirname(dir_param)
        self.full_image_save_dir = os.path.join(os.getcwd(), self.image_save_dir)

        try: # check if directory exists, if not make one
            if not os.path.exists(self.image_save_dir):
                os.makedirs(self.image_save_dir)
            rospy.loginfo("Saving images to: " + self.full_image_save_dir)
        except Exception as e: # if we fail (permissions, etc.)
            rospy.logerr("Failed to create directory: " + e)
            exit()

        rospy.loginfo("Ready to save image from topic: " + self.camera_color_topic)
        rospy.Service('snap_picture', SnapPicture, self.snap_picture)
        while (self.image == None and not rospy.is_shutdown()): pass

    ### @brief ROS Subscriber Callback to get the latest color image
    ### @param msg - ROS Image message
    def camera_color_cb(self, msg):
        with self.img_mutex:
            self.image = msg

    ### @brief ROS Service Callback to save the latest rgb picture with the desired name
    ### @param req - ROS 'SnapPicture' Service message
    def snap_picture(self, req):
        res = SnapPictureResponse()
        res.success = False
        res.filepath = "NULL"
        bridge = CvBridge()
        with self.img_mutex:
            try:
                cv_image = bridge.imgmsg_to_cv2(self.image, "bgr8")
                cv2.imwrite(os.path.join(self.image_save_dir, req.filename), cv_image)
                filepath = os.path.join(self.full_image_save_dir, req.filename)
                rospy.loginfo("Image saved to " + filepath)
                res.success = True
                res.filepath = filepath
            except Exception as e:
                rospy.logerr("Failed to save image.")
                rospy.logerr(e)
        return res

def main():
    rospy.init_node('picture_snapper')
    PictureSnapper()
    rospy.spin()

if __name__=='__main__':
    main()
